package ffte.core.flat

import spinal.core._
import spinal.lib._

import ffte.types.VecComplex
import ffte.types.FixComplex
import ffte.property.{getFFTIW,getFFTOW,FFTIOConfig}
import ffte.algorithm.FFTGen.{Method,msbFFT}

/**
 * Mixed-Space Block (MSB) FFT implementation
 *
 * This class implements FFT using the Mixed-Space Block decomposition method,
 * which is based on the Cooley-Tukey algorithm. MSB decomposition is the most
 * commonly used FFT algorithm and works for any composite transform size N = N1 * N2.
 *
 * Mathematical Foundation:
 * The MSB method implements the Cooley-Tukey decomposition which breaks an N-point
 * DFT into smaller DFTs by reorganizing the computation into a multi-dimensional
 * structure. For N = N1 * N2, the DFT can be written as:
 *
 * X[k1 + N1*k2] = Σ[n2=0 to N2-1] [ W_N^(k2*n2) * Σ[n1=0 to N1-1] ( x[n1 + N1*n2] * W_N1^(k1*n1) ) ] * W_N2^(k2*n1)
 *
 * Where:
 * - W_N = e^(-j*2π/N) is the N-th root of unity
 * - Inner sum: N1-point DFTs (first stage)
 * - Outer sum: N2-point DFTs with twiddle factors (second stage)
 *
 * Hardware Architecture:
 * - First stage: N2 independent N1-point FFTs (parallel computation)
 * - Twiddle factor multiplication: Complex multiplication with W_N^(k2*n1)
 * - Second stage: N1 independent N2-point FFTs (parallel computation)
 * - Data reordering: Row-major to column-major conversion
 *
 * Performance Characteristics:
 * - Works for any composite transform size
 * - Requires N complex multiplications for twiddle factors
 * - Good balance of parallelism and resource usage
 * - Regular data access patterns with stride addressing
 *
 * @tparam T Method type for sub-FFT implementations (Flat, Mixed, Streamed)
 * @param S Scaling shift amount for fixed-point overflow prevention
 * @param p MSB decomposition structure containing stage information and twiddle factor indices
 */
class msb[T <: Method](val S: Int, val p: msbFFT[T]) extends combineFlatComponents[T] {
    /** Standard I/O interface for flat FFT components */
    val io = iogen

    /**
     * First stage FFT processing block
     *
     * This block implements the first stage of the MSB decomposition (Cooley-Tukey),
     * consisting of N2 independent N1-point FFTs operating in parallel. Each FFT
     * processes a subset of the input data in row-major order.
     *
     * Architecture:
     * - N2 parallel FFT processors of size N1
     * - Input width: dW (global input precision)
     * - Output width: interW (intermediate precision)
     * - All processors share the same clock enable signal
     *
     * Data Organization:
     * Input data is divided into N2 blocks, each containing N1 consecutive samples.
     * This implements the inner sum of the Cooley-Tukey decomposition.
     */
    val f_block = FFTIOConfig(dW, interW) on new Area {
        /**
         * Array of first-stage FFT processors
         *
         * Creates N2 independent FFT processors, each handling N1 points.
         * The processors are named "f0", "f1", ..., "f(N2-1)" for debugging.
         */
        val u = (0 until second.N).toList.map { i => fg(first, s"f$i") }

        /**
         * Delay through first stage
         *
         * All processors have identical delay since they use the same
         * implementation method and transform size.
         */
        val delay = u(0).delay

        /**
         * Clock enable distribution
         *
         * All first-stage processors share the same clock enable signal,
         * ensuring synchronous operation across the parallel array.
         */
        u.foreach { x => x.io.ce := io.ce }
    }

    /**
     * Second stage FFT processing block
     *
     * This block implements the second stage of the MSB decomposition, consisting
     * of N1 independent N2-point FFTs operating in parallel. Each FFT processes
     * the twiddle-multiplied intermediate results.
     *
     * Architecture:
     * - N1 parallel FFT processors of size N2
     * - Input width: interW (intermediate precision from first stage)
     * - Output width: oW (global output precision)
     * - All processors share the same clock enable signal
     *
     * Data Organization:
     * Each processor handles N2 samples corresponding to a specific residue
     * class modulo N1, implementing the outer sum of the Cooley-Tukey decomposition.
     */
    val s_block = FFTIOConfig(interW, oW) on new Area {
        /**
         * Array of second-stage FFT processors
         *
         * Creates N1 independent FFT processors, each handling N2 points.
         * The processors are named "f0", "f1", ..., "f(N1-1)" for debugging.
         */
        val u = (0 until first.N).toList.map { i => sg(second, s"f$i") }

        /**
         * Delay through second stage
         *
         * All processors have identical delay since they use the same
         * implementation method and transform size.
         */
        val delay = u(0).delay

        /**
         * Clock enable distribution
         *
         * All second-stage processors share the same clock enable signal,
         * ensuring synchronous operation across the parallel array.
         */
        u.foreach { x => x.io.ce := io.ce }
    }

    /**
     * Twiddle factor buffer registers
     *
     * Stores the intermediate results after twiddle factor multiplication.
     * These registers break the combinatorial path between the first stage
     * FFT outputs and the second stage FFT inputs, improving timing.
     *
     * Structure: 2D array of complex numbers with dimensions (second.N × first.N)
     * - Index (i,k): corresponds to data for second-stage processor k, sample i
     * - Initial value: All zeros
     * - Width: interW bits (intermediate precision)
     */
    val xf = RegInit(Vec(Vec(FixComplex(interW).zero, first.N), second.N))

    /**
     * Data interconnection and twiddle factor processing
     *
     * Implements the complete data flow of the MSB decomposition including
     * input distribution, twiddle factor multiplication, and output collection.
     * The twiddle factors implement the W_N^(k2*n1) term in the Cooley-Tukey equation.
     *
     * Twiddle Factor Calculation:
     * - Index: -p.p.first.store_tab(k) * p.p.second.load_tab(i) mod N
     * - Sign: Negative exponent implements conjugate for forward FFT
     * - Generation: Using TW() function which provides optimal fixed-point twiddle factors
     */
    for (i <- 0 until second.N) {
        for (k <- 0 until first.N) {
            /**
             * Twiddle factor generation
             *
             * Generates the complex twiddle factor W_N^(-k*n1*n2) for the
             * Cooley-Tukey decomposition. The negative sign indicates the
             * forward transform direction.
             *
             * Index calculation combines first and second stage permutation
             * tables to compute the correct twiddle factor exponent.
             */
            val t = TW(-p.p.first.store_tab(k) * p.p.second.load_tab(i), p.N)

            /**
             * Twiddle multiplication and buffering
             *
             * Multiplies the first stage output by the twiddle factor and
             * stores the result in the buffer register. This implements
             * the critical twiddle factor multiplication step of MSB.
             */
            when(io.ce) {
                xf(i)(k) := t.mul(f_block.u(i).io.q.d(k))
            }

            /**
             * First stage to second stage connection
             *
             * Routes the twiddle-multiplied intermediate data to the
             * appropriate second-stage FFT processor input.
             */
            s_block.u(k).io.d.d(i) := xf(i)(k)

            /**
             * Input to first stage connection
             *
             * Routes global input data to the appropriate first-stage processor
             * based on row-major ordering of the 2D decomposition.
             */
            f_block.u(i).io.d.d(k) := io.d.d(i * first.N + k)

            /**
             * Second stage to output connection
             *
             * Routes the output of second-stage processor k, sample i
             * to the global output using column-major ordering.
             */
            io.q.d(k * second.N + i) := s_block.u(k).io.q.d(i)
        }
    }

    /**
     * Total component delay in clock cycles
     *
     * The total delay includes:
     * - First stage FFT delay (f_block.delay)
     * - Second stage FFT delay (s_block.delay)
     * - Additional 1 cycle for the twiddle factor register buffering
     *
     * The extra cycle accounts for the registered twiddle multiplication
     * which breaks the combinatorial path for better timing.
     *
     * @return Total delay from input to output in clock cycles
     */
    def delay = f_block.delay + s_block.delay + 1
}
